# -*- coding: utf-8 -*-
#
#
# Author: Alex Cai
# Email: cyy0523xc@qq.com
# Created Time: 2025-06-13
from traceback import format_exc
from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

# 定义状态结构，包含消息列表（使用add_messages处理器追加消息）
class State(TypedDict):
    messages: Annotated[list, add_messages]

# 初始化状态图构建器
graph_builder = StateGraph(State)

# 模拟LLM调用（不实际请求，直接返回预设响应）
def chatbot(state: State):
    user_message = state["messages"][-1].content
    # 根据用户输入返回不同模拟响应
    if "你好" in user_message or "hi" in user_message.lower():
        response = "你好！我是LangGraph演示机器人，很高兴为你服务。"
    elif "LangGraph" in user_message:
        response = "LangGraph是一个用于构建状态ful LLM应用的框架，支持工具调用和对话记忆。"
    elif "退出" in user_message or "quit" in user_message.lower():
        response = "再见！希望你喜欢这个演示。"
    else:
        response = "抱歉，我暂时还不理解这个问题。你可以问我关于LangGraph的问题。"

    return {"messages": [{"role": "assistant", "content": response}]}

# 添加聊天节点到图中
graph_builder.add_node("chatbot", chatbot)
# 设置入口和结束点
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)
# 编译图
graph = graph_builder.compile()

# 定义流式输出函数（模拟LLM流式响应）
def stream_graph_updates(user_input: str):
    print("Assistant: ", end="", flush=True)
    response = graph.invoke({"messages": [{"role": "user", "content": user_input}]})
    # 模拟流式输出效果
    assistant_msg = response["messages"][-1].content
    for char in assistant_msg:
        print(char, end="", flush=True)
        import time
        time.sleep(0.05)  # 模拟打字延迟
    print()

# 对话循环
print("===== LangGraph 基础聊天机器人演示 =====")
print("输入'退出'、'quit'或'q'结束对话")
print()

while True:
    try:
        user_input = input("User: ")
        if user_input.lower() in ["退出", "quit", "q"]:
            print("Assistant: 再见！")
            break
        stream_graph_updates(user_input)
    except Exception as e:
        print(f"Assistant: 发生错误: {str(e)}")
        print(format_exc())
        break
